package dns

import (
	"bufio"
	"crypto/rand"
	"dnsmap/api"
	"encoding/csv"
	"encoding/hex"
	"fmt"
	"github.com/bobesa/go-domain-util/domainutil"
	"github.com/dean2021/hackpool"
	"github.com/fatih/color"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
	"log"
	"net"
	"os"
	"reflect"
	"strings"
	"time"
)

var (
	wildcard []string
	Records  []Record
	g        = color.New(color.FgGreen)
)

type DNSMap struct{}

type Record struct {
	Domain string
	IP     []string
}

func New(c *cli.Context) error {

	domain := c.String("domain")
	concurrency := c.Int("concurrency")
	timeout := time.Duration(c.Int64("timeout")) * time.Second
	output := c.String("output")
	wordList := c.String("wordlist")
	debug := c.Bool("debug")
	start := time.Now()

	if domain == "" {
		_ = cli.ShowAppHelp(c)
		return errors.New("--domain 不能为空")
	}

	if !debug {
		fmt.Println("debug.")
		fmt.Println("timeout", timeout)
		fmt.Println("output", output)
	}

	// 泛解析检查
	baseDomain := domainutil.Domain(domain)
	hasWildcard := false
	hasWildcard, wildcard, _ = detectWildcard(baseDomain)
	if hasWildcard {
		fmt.Println("泛解析IP:", wildcard)
	}

	taskChan := mixInAPIDict(domain, wordList, timeout)
	hp := hackpool.New(concurrency, Worker)
	go func() {
		for t := range taskChan {
			hp.Push(t)
		}
		hp.Close()
	}()

	hp.Run()

	// 输出
	file, err := os.Create(output)
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	// csv
	csvOut := csv.NewWriter(file)
	defer csvOut.Flush()
	_ = csvOut.Write([]string{"Domain", "IP"})

	counter := 0
	for _, record := range Records {
		counter++
		_ = csvOut.Write([]string{record.Domain, strings.Join(record.IP, ",")})
	}

	log.Printf("done in %.2f seconds, %d records\n", time.Since(start).Seconds(), counter)
	return nil
}

// Lookup a random host to determine if a wildcard A record exists
// Adapted from https://github.com/jrozner/sonar/blob/master/wildcard.go
func detectWildcard(domain string) (bool, []string, error) {
	bytes := make([]byte, 16)
	_, err := rand.Read(bytes)
	if err != nil {
		return false, nil, err
	}
	domain = fmt.Sprintf("%s.%s", hex.EncodeToString(bytes), domain)
	fmt.Println("检查泛解析:", domain)
	answers, err := net.LookupHost(domain)
	if err != nil {
		if asserted, ok := err.(*net.DNSError); ok && asserted.Err == "no such host" {
			return false, nil, nil
		}
		return false, nil, err
	}
	return true, answers, nil
}

func IsPublicIP(IP net.IP) bool {
	if IP.IsLoopback() || IP.IsLinkLocalMulticast() || IP.IsLinkLocalUnicast() {
		return false
	}
	if ip4 := IP.To4(); ip4 != nil {
		switch true {
		case ip4[0] == 10:
			return false
		case ip4[0] == 172 && ip4[1] >= 16 && ip4[1] <= 31:
			return false
		case ip4[0] == 192 && ip4[1] == 168:
			return false
		case ip4[0] == 0 && ip4[1] == 0:
			return false
		default:
			return true
		}
	}
	return false
}

func mixInAPIDict(domain, dict string, timeout time.Duration) <-chan string {
	subDomainsToQuery := make(chan string)
	mix := make(chan string)

	// mix in
	go func() {
		defer close(subDomainsToQuery)

		domains := map[string]struct{}{}
		for sub := range mix {
			domains[sub] = struct{}{}
		}

		for domain := range domains {
			subDomainsToQuery <- domain
		}
	}()

	// API
	// 同步，保证 API 完整执行
	for sub := range api.Query(domain, timeout) {
		mix <- sub
	}

	go func() {
		defer close(mix)

		// Domain
		mix <- domain

		// Dict
		file, err := os.Open(dict)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			mix <- scanner.Text() + "." + domain
		}
		if err := scanner.Err(); err != nil {
			log.Fatal(err)
		}
	}()

	return subDomainsToQuery
}

func Worker(i interface{}) {
	host := i.(string)
	if addrs, err := net.LookupHost(host); err == nil {
		// 忽略泛解析域名
		if reflect.DeepEqual(addrs, wildcard) {
			return
		}

		// 仅保留公网域名
		if IsPublicIP(net.ParseIP(addrs[0])) {
			_, _ = g.Printf("%25s", host)
			fmt.Printf(" : %v \n", addrs)
			Records = append(Records, Record{Domain: host, IP: addrs})
		}
	}
}
